#!/usr/bin/env python3

import argparse
import grpc
import time
import sys
import os 
from p4.v1 import p4runtime_pb2
from p4.v1 import p4runtime_pb2_grpc
from p4.config.v1 import p4info_pb2
import google.protobuf.text_format
import socket
import struct

def get_table_id(p4info, name):
    """获取表的ID"""
    for table in p4info.tables:
        if table.preamble.name == name:
            return table.preamble.id
    raise ValueError(f"Table {name} not found")

def get_action_id(p4info, name):
    """获取动作的ID"""
    for action in p4info.actions:
        if action.preamble.name == name:
            return action.preamble.id
    raise ValueError(f"Action {name} not found")

def get_field_id(p4info, table_name, field_name):
    """获取表字段的ID"""
    for table in p4info.tables:
        if table.preamble.name == table_name:
            for field in table.match_fields:
                if field.name == field_name:
                    return field.id
    raise ValueError(f"Field {field_name} not found in table {table_name}")

def get_param_id(p4info, action_name, param_name):
    """获取动作参数的ID"""
    for action in p4info.actions:
        if action.preamble.name == action_name:
            for param in action.params:
                if param.name == param_name:
                    return param.id
    raise ValueError(f"Param {param_name} not found in action {action_name}")

def set_ecn_threshold(switch_ip, p4info_path, threshold):
    """设置ECN阈值"""
    
    # 创建gRPC通道
    channel = grpc.insecure_channel(switch_ip)
    stub = p4runtime_pb2_grpc.P4RuntimeStub(channel)
    
    # 读取P4Info文件
    p4info = p4info_pb2.P4Info()
    with open(p4info_path, 'r') as f:
        google.protobuf.text_format.Merge(f.read(), p4info)
    
    # 设置选举ID
    election_id = p4runtime_pb2.Uint128(high=0, low=1)
    
    try:
        # 开启流通道
        stream = stub.StreamChannel()
        
        # 发送仲裁消息
        request = p4runtime_pb2.StreamMessageRequest()
        request.arbitration.device_id = 1
        request.arbitration.election_id.CopyFrom(election_id)
        stream.send(request)
        
        # 等待仲裁完成
        response = next(stream)
        if not response.HasField('arbitration'):
            print("Arbitration failed")
            return False
            
        print(f"Arbitration successful, setting ECN threshold to {threshold}")
        
        # 创建表项来设置ECN阈值
        table_entry = p4runtime_pb2.TableEntry()
        table_entry.table_id = get_table_id(p4info, "Ingress.ecn_threshold_config")
        
        # 匹配所有IPv4数据包 (0.0.0.0/0)
        match = table_entry.match.add()
        match.field_id = get_field_id(p4info, "Ingress.ecn_threshold_config", "hdr.ipv4.dstAddr")
        match.lpm.prefix_len = 0
        match.lpm.value = b'\x00\x00\x00\x00'  # 0.0.0.0
        
        # 设置动作和参数
        action = table_entry.action.action
        action.action_id = get_action_id(p4info, "Ingress.set_ecn_threshold")
        
        # 设置阈值参数
        param = action.params.add()
        param.param_id = get_param_id(p4info, "Ingress.set_ecn_threshold", "threshold")
        param.value = threshold.to_bytes(4, byteorder='big')
        
        # 创建写入请求
        update = p4runtime_pb2.Update()
        update.type = p4runtime_pb2.Update.INSERT
        update.entity.table_entry.CopyFrom(table_entry)
        
        write_request = p4runtime_pb2.WriteRequest()
        write_request.device_id = 1
        write_request.election_id.CopyFrom(election_id)
        write_request.updates.append(update)
        
        # 发送写入请求
        stub.Write(write_request)
        
        print(f"Successfully configured ECN threshold: {threshold}")
        return True
        
    except Exception as e:
        print(f"Error setting ECN threshold: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='ECN Dynamic Threshold Controller')
    parser.add_argument('--switch-ip', default='localhost:50051', 
                       help='Switch IP and port (default: localhost:50051)')
    parser.add_argument('--p4info', default='build/ecn.p4info.txt', 
                       help='P4Info file path (default: build/ecn.p4info.txt)')
    parser.add_argument('--threshold', type=int, required=True, 
                       help='ECN threshold value (queue depth)')
    
    args = parser.parse_args()
    
    # 检查P4Info文件是否存在
    if not os.path.exists(args.p4info):
        print(f"P4Info file not found: {args.p4info}")
        print("Please compile the P4 program first with 'make'")
        sys.exit(1)
    
    # 设置ECN阈值
    success = set_ecn_threshold(args.switch_ip, args.p4info, args.threshold)
    
    if success:
        print("ECN configuration completed successfully!")
    else:
        print("ECN configuration failed!")
        sys.exit(1)

if __name__ == '__main__':
    main()